#include <iostream>
#include <fstream>
#include <vector>
#include "txtpara.h"
#include "KDindex.h"
#include "directsampling.h"
#include "taxonomy_tree.h"
#include "tplate.h"
#include "bayesian.h"

using namespace std;

int main(int argc, char *argv[])
{
    if (argc < 6)
    {
        cout << "\nUsage:\n" << argv[0] <<
            " param.cf ti.sgems model.txt template.par result\n\n";
        exit(0);
    }
    /* direct sampling parameters */
    int NR_neighbor;            //Maximum number of neighbors
    int NR_try;                    //Maximum number of search
    float tolerance;            //tolerance of difference
    int xsize, ysize, zsize;
    float unitsize = 1.0;
    char *config = argv[1];
    int search_radius = 100;
    NR_neighbor = getparaINT("NR_neighbor", config);
    tolerance = getparaDOUBLE("tolerance", config);
    xsize = getparaDOUBLE("xsize", config);
    ysize = getparaDOUBLE("ysize", config);
    zsize = getparaDOUBLE("zsize", config);
    double percentage = getparaFLOAT("percent", config);
    std::vector < coord > sample;
    //Step 1.  Generate sample points for the first time
    //>>>>>1.1 calculating Number of points
    int tol_sample = xsize * ysize * zsize * percentage;
    int N_sample = xsize * ysize * zsize * percentage;
    int interval = round(sqrt(xsize * ysize / N_sample));
    if ((xsize / (interval)) * (ysize / (interval)) > N_sample)
        interval++;
    //>>>>>1.2 generate random path
    srand(clock());
    std::vector < int >ids;
    int zpos = 0;
    int s_no = 0;
    for (int xx = 0; xx < xsize; xx += interval)
    {
        for (int yy = 0; yy < ysize; yy += interval)
        {
            coord s;
            s.x = xx;
            s.y = yy;
            s.z = zpos;
            sample.push_back(s);
            ids.push_back(s_no);
            s_no++;
        }
    }
    N_sample = sample.size();
    for (int xpos = 0; xpos < xsize; xpos++)
    {
        for (int ypos = 0; ypos < ysize; ypos++)
        {
            if (((xpos % interval) == 0) && ((ypos % interval) == 0))
                continue;
            coord s;
            s.x = xpos;
            s.y = ypos;
            s.z = zpos;
            sample.push_back(s);
            ids.push_back(s_no);
            s_no++;
        }
    }
    for (int i = N_sample; i < tol_sample; i++)
    {
        int ranid = N_sample + rand() % (xsize * ysize - N_sample);
        int tmp = ids[i];
        ids[i] = ids[ranid];
        ids[ranid] = tmp;
    }
    for (s_no = 0; s_no < N_sample; s_no++)
    {
        int ranid = rand() % N_sample;
        int tmp = ids[s_no];
        ids[s_no] = ids[ranid];
        ids[ranid] = tmp;
    }
    N_sample = tol_sample;
    //Step 2.  The simulation on the coarsest grid
    Index < L2 < float >>*R = createIndex();
    float *ti;
    int tilength, tiwidth, tiheight;
    int start_id = 0;
    int level = 0;
    int *target = new int[int (xsize * ysize * zsize)];
    memset(target, -1, sizeof(int) * xsize * ysize * zsize);
    load_ti_3D(&ti, argv[2], &tilength, &tiwidth, &tiheight);
    NR_try =
        int (getparaDOUBLE("Fraction", config) * tilength * tiwidth *
             tiheight);
    Simulation(R, ti, NR_neighbor, NR_try, tolerance, sample, ids, target,
               N_sample, tilength, tiwidth, tiheight, unitsize, start_id,
               level, xsize, ysize, zsize, search_radius);
    //Step 3. Interpolation remaining parts

    int length = xsize, width = ysize;
    int NR_cate;
    float *marginal_prob;
    float *marginal_prob_target;
    int maxdv = getparaINT("maxdv", config);
    NR_cate = getparaINT("NR_CATE", config);
    marginal_prob = new float[NR_cate];
    marginal_prob_target = new float[NR_cate];
    getparaFLOATARRAY("marginal_prob", config, marginal_prob, NR_cate);
    getparaFLOATARRAY("target_marginal_prob", config, marginal_prob_target,
                      NR_cate);
    float *logmarginal = new float[NR_cate];
    float *rectmarginal = new float[NR_cate];
    for (int c = 0; c < NR_cate; c++)
    {
        logmarginal[c] = log(marginal_prob[c]);
    }
    // Read the training image to the trianing grid
    tplate tp;
    load_template(&tp, argv[4]);
    ALL_taxonomy_tree *all_tree = new ALL_taxonomy_tree;
    load_taxonomy_tree(all_tree, tp.tpsize, argv[3]);
    train_grid *ti_grid;
    target_grid *grid = new target_grid;
    init_sim_grid(grid, length, width, 1, NR_cate, N_sample, target);
    int *hard_cate = new int[NR_cate];
    int g = 0;
    float *prob_result = new float[NR_cate];
    int cur_length = grid->length[g];
    int cur_width = grid->width[g];
    float *marginal = new float[NR_cate];
    float *weight = new float[2];
    for (int x = 0; x < cur_length; x++)
    {
        for (int y = 0; y < cur_width; y++)
        {
            memcpy(hard_cate,
                   grid->sample[g].category +
                   x * cur_width * NR_cate + y * NR_cate,
                   NR_cate * sizeof(int));
            data_event *dv = new data_event;
            dv->tpos = new int[tp.tpsize];
            dv->pos = new int[tp.tpsize * 2];
            dv->value = new int *[tp.tpsize];
            get_dv(dv, &tp, grid, g, x, y, maxdv);
            if (dv->dvsize <= 0)
            { 
                memcpy(grid->sample[g].prob +
                       x * cur_width * NR_cate + y * NR_cate,
                       marginal_prob, sizeof(float) * NR_cate);
            }
            else
            {
                marginal = marginal_prob_target;
                memset(prob_result, 0, sizeof(float) * NR_cate);
                float *all_node_prob =
                    extract_all_pdf(dv, all_tree->t_tree[g],
                                    marginal, NR_cate);
                for (int c = 0; c < NR_cate; c++)
                {
                    rectmarginal[c] = log(marginal[c]);
                }
                for (int ii = 0; ii < dv->dvsize; ii++)
                {
                    float sum = 0;
                    for (int c = 0; c < NR_cate; c++)
                    {
                        sum += all_node_prob[ii * 2 + c];
                    }
                    for (int c = 0; c < NR_cate; c++)
                    {
                        all_node_prob[ii * 2 + c] /= sum;
                    }
                }
                merge_prob(prob_result, all_node_prob, logmarginal,
                           rectmarginal, NR_cate, dv->dvsize);
                memcpy(grid->sample[g].prob + x * cur_width * NR_cate +
                       y * NR_cate, prob_result, sizeof(float) * NR_cate);
            }
            for (int i = 0; i < dv->dvsize; i++)
            {
                delete dv->value[i];
            }
            delete dv->value;
            delete dv->tpos;
            delete dv->pos;
            delete dv;
        }
    }

    char tmpstr[256];
    sprintf(tmpstr, "%s.txt", argv[5]);
    g = 0;
    ofstream ofs(tmpstr);
    for (int y = 0; y < grid->width[g]; y++)
    {
        for (int x = 0; x < grid->length[g]; x++)
        {
            double max = 0;
            int which = -1;
            for (int c = 0; c < NR_cate; c++)
            {
                float cate_prob =
                    grid->sample[g].prob[x * grid->width[g] * NR_cate +
                                         y * NR_cate + c];
                if (cate_prob > max)
                {
                    max = cate_prob;
                    which = c;
                }
            }
            ofs << which << " ";
        }
        ofs << endl;
    }
    ofs.close();
    delete[]hard_cate;
    return 0;
}
